import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics, callbacks


class KerasMobileNetAgain(keras.Model):

    def __init__(self, input_shape, n_cls, **kwargs):
        super().__init__(**kwargs)
        self.base_model = keras.applications.MobileNet(
            input_shape=input_shape,
            include_top=False,
            weights='imagenet',
            pooling='avg',
        )
        self.base_model.trainable = False
        self.customer_model = layers.Dense(n_cls, activation=activations.softmax)  # ATTENTION softmax is fatal!

    def call(self, x, training=None):
        x = self.base_model(x, training=training)
        x = self.customer_model(x, training=training)
        return x


if '__main__' == __name__:

    import python_ai.common.read_data.cat_dog_teachers as cat_dog
    import os
    from KerasMobileNet import get_model

    VER = 'v4.1'
    ALPHA = 1e-5
    SIZE = 224
    BATCH_SIZE = 8
    N_EPOCH = 4
    BASE_DIR, FILE_NAME = os.path.split(__file__)
    dir = '../../../../../large_data/DL1/_many_files/zoo'
    IMG_DIR = os.path.join(BASE_DIR, dir)
    LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)

    (x_train, y_train), (x_val, y_val), (x_test, y_test) = cat_dog.load_data(IMG_DIR, 1, SIZE, SIZE, ch_first=False)

    input_shape = (SIZE, SIZE, 3)
    model = KerasMobileNetAgain(input_shape, 2)
    model.build(input_shape=[None, *input_shape])
    # model = get_model(2, input_shape)
    model.summary()

    model.compile(
        loss=losses.sparse_categorical_crossentropy,
        optimizer=optimizers.Adam(learning_rate=ALPHA),
        metrics=[metrics.sparse_categorical_accuracy]
    )

    model.fit(
        x_train,
        y_train,
        BATCH_SIZE,
        N_EPOCH,
        callbacks=[callbacks.TensorBoard(LOG_DIR, update_freq='batch', profile_batch=0)],
        validation_data=(x_val, y_val),
        validation_batch_size=BATCH_SIZE,
    )

    print('Training...')
    model.evaluate(x_test, y_test, BATCH_SIZE)
